import torch
from torch.nn import Module, Linear, ModuleList, LayerNorm, Sequential, GELU
from seasoncnnblock import CnnSeasonBlock

class CnnSeason(Module):
    def __init__(self, layer:int, topk:int, model_dim:int, seq:int):
        super(CnnSeason, self).__init__()
        self.topk = topk
        self.embed = Linear(model_dim, model_dim)
        self.layer = layer
        self.models = ModuleList([CnnSeasonBlock(self.topk, model_dim) for _ in range(layer)])
        self.layer_norm = LayerNorm(model_dim)
        self.dembed = Linear(model_dim, 1)
        self.perdict = Sequential(Linear(seq, seq), GELU(), Linear(seq, 1))


    def forward(self, season:torch.Tensor):
        # 先norm
        means = season.mean(dim=-2, keepdim=True).detach()
        season_level = season - means
        stdev = torch.sqrt(torch.var(season, dim=1, keepdim=True, unbiased=False) + 1e-5)
        season_normed = season_level / stdev

        # 对season进行特征embedding，将特征数从input_features变为model_dim
        embeded = self.embed(season_normed)

        # cnnseasonblock层
        for i in range(self.layer):
            embeded = self.layer_norm(self.models[i](embeded))
        
        # denorm，用目标特征的均值方差来做
        output = embeded * stdev + means
        return output


if __name__ == "__main__":
    B = 2
    N = 5
    T = 13
    C = 4
    x = torch.ones(B, T, C)
    block = CnnSeason(layer=3, topk=3, model_dim=C, seq=T)
    res = block(x)